#pragma once
#include "dbghelp.h"
#pragma comment( lib, "dbghelp.lib" )

void		OutputInfo(LPCTSTR lpszText);
void		OutputInfo(LPSTR lpszText);

class ProcedureAddrRetrieve
{
public:
	ProcedureAddrRetrieve(TCHAR *PdbSearchPath,DWORD _Options = NULL);
	~ProcedureAddrRetrieve();

public:
	BOOL		LoadSymbol(TCHAR *DllName);
	BOOL		EnumSymbol(TCHAR *pSearchMask = NULL);
	DWORD64		RetrieveAddr(TCHAR *_szProcedureName);

private:
	BOOL		GetFileSize( const TCHAR* pFileName, DWORD& FileSize );
	static BOOL CALLBACK EnumSymbolsCallback( SYMBOL_INFO* pSymInfo, ULONG SymbolSize, PVOID UserContext );
	
private:
	DWORD64		ModBase64;
	char		szProcedureName[MAX_PATH];
	DWORD64		Address;
};

ProcedureAddrRetrieve::ProcedureAddrRetrieve(TCHAR *PdbSearchPath,DWORD _Options )
{
	BOOL bRet = FALSE;
	ZeroMemory(this,sizeof(ProcedureAddrRetrieve));
	// Set options
	DWORD Options = SymGetOptions();

	// SYMOPT_DEBUG option asks DbgHelp to print additional troubleshooting
	// messages to debug output - use the debugger's Debug Output window
	// to view the messages

	Options =Options | SYMOPT_DEBUG | _Options/*|SYMOPT_ALLOW_ABSOLUTE_SYMBOLS*/;

	SymSetOptions( Options );
	// Initialize DbgHelp and load symbols for all modules of the current process

	bRet = SymInitializeW(
		GetCurrentProcess(),  // Process handle of the current process
		PdbSearchPath,                 // user-defined search path -> use default
		FALSE                 // Do not load symbols for modules in the current process
		);

	if( !bRet )
	{
		OutputInfo(L"Error: SymInitialize() failed. Error code: %u \n");
		return ;
	}
}
BOOL ProcedureAddrRetrieve::LoadSymbol(TCHAR *DllName)
{
	HMODULE	ModBase;
	TCHAR	DllFullPath[MAX_PATH];
	DWORD	FileSize  = 0;
	char	lpszFullPath[MAX_PATH];

	//Get DLL full path.
	ModBase = GetModuleHandle(DllName);
	if(ModBase == NULL)
	{
		ModBase = LoadLibrary(DllName);
		if(ModBase == NULL)
			return FALSE;
	}
	GetModuleFileName(ModBase,DllFullPath,MAX_PATH);

	GetFileSize(DllFullPath, FileSize);

	int nLen = WideCharToMultiByte(CP_ACP, 0, DllFullPath, -1,
		NULL, 0, NULL, NULL);
	WideCharToMultiByte(CP_ACP, 0, DllFullPath, -1, lpszFullPath, nLen,
		NULL, NULL);

	if(ModBase64)
	{
		// Unload symbols for the module
		if(!SymUnloadModule64( GetCurrentProcess(), ModBase64 ))
		{
			OutputInfo(L"Error: SymUnloadModule64() failed. Error code: %u \n");
			return FALSE;
		}
	}

	ModBase64 = SymLoadModule64(
		GetCurrentProcess(),		// Process handle of the current process
		NULL,						// Handle to the module's image file (not needed)
		lpszFullPath,				// Path/name of the file
		NULL,						// User-defined short name of the module (it can be NULL)
		(DWORD64)ModBase,			// Base address of the module (cannot be NULL if .PDB file is used, otherwise it can be NULL)
		FileSize					// Size of the file (cannot be NULL if .PDB file is used, otherwise it can be NULL)
		);

	if( ModBase == 0 )
	{
		OutputInfo(L"Error: SymLoadModule64() failed. Error code: ");
		return FALSE;
	}
	return TRUE;
}
BOOL ProcedureAddrRetrieve::GetFileSize( const TCHAR* pFileName, DWORD& FileSize )
{
	// Check parameters
	if( pFileName == 0 )
	{
		return FALSE;
	}

	// Open the file
	HANDLE hFile = CreateFile( pFileName, GENERIC_READ, FILE_SHARE_READ,
		NULL, OPEN_EXISTING, 0, NULL );

	if( hFile == INVALID_HANDLE_VALUE )
	{
		OutputInfo(L"CreateFile() failed. Error:");
		return FALSE;
	}

	// Obtain the size of the file
	FileSize = ::GetFileSize( hFile, NULL );
	if( FileSize == INVALID_FILE_SIZE )
	{
		OutputInfo(L"GetFileSize() failed. Error: ");
		// and continue ...
	}

	// Close the file
	if( !CloseHandle( hFile ) )
	{
		OutputInfo(L"CloseHandle() failed. Error: ");
		// and continue ...
	}
	// Complete

	return ( FileSize != INVALID_FILE_SIZE );

}
ProcedureAddrRetrieve::~ProcedureAddrRetrieve()
{
	if(ModBase64)
	{
		// Unload symbols for the module
		if(!SymUnloadModule64( GetCurrentProcess(), ModBase64 ))
		{
			OutputInfo(L"Error: SymUnloadModule64() failed. Error code: ");
		}
	}
}
BOOL ProcedureAddrRetrieve::EnumSymbol(TCHAR *pSearchMask)
{
	BOOL bRet = TRUE;
	char	lpszMask[1024];

	// Enumerate symbols and display information about them
	int nLen = WideCharToMultiByte(CP_ACP, 0, pSearchMask, -1,
		NULL, 0, NULL, NULL);
	WideCharToMultiByte(CP_ACP, 0, pSearchMask, -1, lpszMask, nLen,
		NULL, NULL);

	bRet = ::SymEnumSymbols(
		GetCurrentProcess(),	// Process handle of the current process
		ModBase64,				// Base address of the module
		lpszMask,				// Mask (NULL -> all symbols)
		EnumSymbolsCallback,	// The callback function
		this					// A used-defined context can be passed here, if necessary
		);

	if( !bRet )
	{
		OutputInfo(L"Error: SymEnumSymbols() failed. Error code: ");
	}
	return TRUE;
}
BOOL CALLBACK ProcedureAddrRetrieve::EnumSymbolsCallback( SYMBOL_INFO* pSymInfo, ULONG SymbolSize, PVOID UserContext )
{
	ProcedureAddrRetrieve *pPar = (ProcedureAddrRetrieve *)UserContext;
	if( pSymInfo != 0 )
	{
		if(*pPar->szProcedureName)
		{
			if(strcmp(pPar->szProcedureName,pSymInfo->Name)==0)
			{
				pPar->Address = pSymInfo->Address;
				return FALSE;
			}
		}
		else
		{
			OutputInfo(pSymInfo->Name);
		}
	}
	return TRUE;
}
DWORD64 ProcedureAddrRetrieve::RetrieveAddr(TCHAR *_szProcedureName)
{
	//strncpy(szProcedureName,_szProcedureName,MAX_PATH);
	Address = 0;
	if(!EnumSymbol(L""))
		return NULL;
	return Address;
}

TCHAR lpszAll[100];
TCHAR lpszOut[200];

void OutputInfo(LPCTSTR lpszText)
{
	swprintf(lpszAll, L"%ws %d", lpszText, ::GetLastError());
}

void OutputInfo(LPSTR lpszText)
{
	int nLen = MultiByteToWideChar(CP_ACP, 0, lpszText, -1,
		NULL, 0);
	MultiByteToWideChar(CP_ACP, 0, lpszText, -1, lpszAll, nLen);

	swprintf(lpszOut, L"%ws %d", lpszAll, ::GetLastError());
}